# 导入必要的库和模块
from pathlib import Path
import pickle

import gradio as gr
import numpy as np
from PIL import Image, ImageFilter
from typing import Any


# 加载保存的KNN模型
MODEL_PATH = Path(__file__).with_name("best_knn_model.pkl")
if not MODEL_PATH.exists():
    raise FileNotFoundError(
        f"未找到已训练模型 {MODEL_PATH}. 请先运行 `optimal_knn.py` 生成模型。"
    )
with MODEL_PATH.open("rb") as f:
    model = pickle.load(f)


def to_pil_image(image: Any) -> Image.Image:
    """将各种可能的输入（dict/ndarray/PIL/路径）统一转为 PIL.Image。

    兼容 Gradio v5 中 `Sketchpad`/`Image` 可能返回的 dict 结构。
    同时处理透明通道，合成白底。
    """
    if image is None:
        raise ValueError("未收到图像，请先在画板上绘制或上传图片。")

    if isinstance(image, Image.Image):
        img = image
    elif isinstance(image, np.ndarray):
        arr = image
        if arr.dtype != np.uint8:
            arr = np.clip(arr, 0, 255).astype(np.uint8)
        img = Image.fromarray(arr)
    elif isinstance(image, dict):
        # Gradio v5 的一些组件会返回包含图像的字典
        for key in ("image", "composite", "background"):
            if key in image and image[key] is not None:
                return to_pil_image(image[key])
        path_val = image.get("path") or image.get("image_path")
        if isinstance(path_val, str):
            p = Path(path_val)
            if p.exists():
                return Image.open(p)
        raise TypeError("无法从字典数据中提取图像。")
    elif isinstance(image, str):
        p = Path(image)
        if p.exists():
            return Image.open(p)
        raise TypeError("提供的字符串不是有效的图像路径。")
    else:
        raise TypeError(f"不支持的图像类型: {type(image)}")

    # 统一去除透明通道（若有），合成白底
    if img.mode in ("RGBA", "LA"):
        base = Image.new("RGBA", img.size, (255, 255, 255, 255))
        img = Image.alpha_composite(base, img.convert("RGBA")).convert("RGB")
    return img


def preprocess_image_to_vector(img: Image.Image, deskew: bool = True) -> np.ndarray:
    """把用户上传/绘制的图片转成与 digits 数据集一致的 8x8=64 维特征。

    - 转为灰度
    - 提取前景（自适应反白）并裁剪到前景边界
    - 居中填充为正方形后，调整为 8x8
    - 反白（使前景为高值）以贴近 digits 分布
    - 归一化到 [0, 16]（digits 原数据大致在 0-16）
    - 展平为 (1, 64)
    """
    # 转灰度
    img_gray = img.convert("L")

    # to numpy float32
    gray = np.asarray(img_gray, dtype=np.float32)

    # 自适应前景朝向：若整体偏亮（白底深色笔迹），则反转使前景为高值
    if gray.mean() > 127:
        gray = 255.0 - gray

    # 提取前景掩码，用较低阈值保留弱笔画（提升对“7”的召回）
    thresh = max(8.0, gray.max() * 0.15)
    mask = gray > thresh

    # 裁剪到前景最小外接矩形，并保留边距（使外观更像 digits）
    if mask.any():
        ys, xs = np.where(mask)
        y0, y1 = int(ys.min()), int(ys.max()) + 1
        x0, x1 = int(xs.min()), int(xs.max()) + 1
        cropped = gray[y0:y1, x0:x1]
    else:
        cropped = gray

    # 增加 20% 边距并居中到正方形画布
    h, w = cropped.shape
    side = int(np.ceil(max(h, w) * 1.2))
    side = max(side, 1)
    square = np.zeros((side, side), dtype=np.float32)
    y_off = (side - h) // 2
    x_off = (side - w) // 2
    square[y_off:y_off + h, x_off:x_off + w] = cropped

    # 基于二阶矩的去倾斜（deskew），仅在倾角显著时生效
    if deskew:
        w_sum = float(square.sum())
        if w_sum > 1e-3:
            ys, xs = np.indices(square.shape)
            cy = float((ys * square).sum() / w_sum)
            cx = float((xs * square).sum() / w_sum)
            y = ys - cy
            x = xs - cx
            cov_xx = float((square * (x * x)).sum() / w_sum)
            cov_yy = float((square * (y * y)).sum() / w_sum)
            cov_xy = float((square * (x * y)).sum() / w_sum)
            angle = 0.5 * np.arctan2(2.0 * cov_xy, (cov_xx - cov_yy) + 1e-12)
            angle_deg = float(np.degrees(angle))
            if abs(angle_deg) > 6.0:
                pil_sq = Image.fromarray(np.clip(square, 0, 255).astype(np.uint8))
                try:
                    rotated = pil_sq.rotate(-angle_deg, resample=Image.BILINEAR, expand=True, fillcolor=0)
                except TypeError:
                    rotated = pil_sq.rotate(-angle_deg, resample=Image.BILINEAR, expand=True)
                    bg = Image.new("L", rotated.size, 0)
                    bg.paste(rotated, (0, 0))
                    rotated = bg
                square = np.asarray(rotated, dtype=np.float32)

    # 轻微形态学膨胀以加粗细笔画，再轻微高斯模糊
    square_img = Image.fromarray(np.clip(square, 0, 255).astype(np.uint8)).filter(
        ImageFilter.MaxFilter(size=3)
    ).filter(ImageFilter.GaussianBlur(radius=0.6))

    # 先平均下采样到 16x16，再做 2x2 平均池化到 8x8，保留结构细节
    try:
        resample = Image.Resampling.BOX  # 平均法
    except Exception:
        resample = Image.BOX
    medium_img = square_img.resize((16, 16), resample=resample)
    medium = np.asarray(medium_img, dtype=np.float32)
    small = medium.reshape(8, 2, 8, 2).mean(axis=(1, 3))

    data = small.astype(np.float32)

    # 直接线性映射到 0-16（避免每张图独立拉伸导致分布漂移）
    data = (data / 255.0) * 16.0

    vector = data.reshape(1, -1)
    return vector


# 定义预测函数（供 Gradio 调用）
def predict(image: Any) -> str:
    pil_image = to_pil_image(image)

    # Test-Time Augmentation: small rotations to improve robustness (esp. '7')
    angles = [-12, -8, -4, 0, 4, 8, 12]

    def safe_rotate_rgb(img: Image.Image, angle: float) -> Image.Image:
        try:
            return img.rotate(angle, resample=Image.BILINEAR, expand=False, fillcolor=(255, 255, 255))
        except TypeError:
            rotated = img.rotate(angle, resample=Image.BILINEAR, expand=False)
            bg = Image.new("RGB", img.size, (255, 255, 255))
            bg.paste(rotated, (0, 0))
            return bg

    vectors: list[np.ndarray] = []
    for a in angles:
        img_rot = safe_rotate_rgb(pil_image, a)
        vec = preprocess_image_to_vector(img_rot, deskew=False)
        vectors.append(vec)

    X_aug = np.vstack(vectors)

    # 预测与概率（聚合 TTA 概率）
    try:
        if hasattr(model, "predict_proba"):
            proba_all = model.predict_proba(X_aug)
            proba_mean = proba_all.mean(axis=0)
            label = int(np.argmax(proba_mean))
        else:
            preds = model.predict(X_aug)
            # 多数投票
            counts = np.bincount(preds.astype(int), minlength=10)
            label = int(np.argmax(counts))
    except Exception as e:
        return f"预测失败: {e}"

    return f"预测: {label}"


# 创建 Gradio 接口
with gr.Blocks(title="Optimal KNN - Digits") as demo:
    gr.Markdown("**手写数字识别（KNN）** - 上传或在画板上手写一个数字。")

    with gr.Row():
        sketch = gr.Sketchpad(
            label="image",
            brush=gr.Brush(colors=["#000000"], default_size=18),
            height=256,
            width=256,
        )
        output = gr.Textbox(label="output", interactive=False)
    with gr.Row():
        clear_btn = gr.ClearButton([sketch, output], value="Clear")
        submit_btn = gr.Button("Submit", variant="primary")

    submit_btn.click(predict, inputs=sketch, outputs=output)


# 启动接口
if __name__ == "__main__":
    demo.launch()

